{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "F7aJhsgLAWvO"
   },
   "source": [
    "# Style-Bert-VITS2 (ver 2.5.0) のGoogle Colabでの学習\n",
    "\n",
    "Google Colab上でStyle-Bert-VITS2の学習を行うことができます。\n",
    "\n",
    "このnotebookでは、通常使用ではあなたのGoogle Driveにフォルダ`Style-Bert-VITS2`を作り、その内部での作業を行います。他のフォルダには触れません。\n",
    "Google Driveを使わない場合は、初期設定のところで適切なパスを指定してください。\n",
    "\n",
    "## 流れ\n",
    "\n",
    "### 学習を最初からやりたいとき\n",
    "上から順に実行していけばいいです。音声合成に必要なファイルはGoogle Driveの`Style-Bert-VITS2/model_assets/`に保存されます。また、途中経過も`Style-Bert-VITS2/Data/`に保存されるので、学習を中断したり、途中から再開することもできます。\n",
    "\n",
    "### 学習を途中から再開したいとき\n",
    "0と1を行い、3の前処理は飛ばして、4から始めてください。スタイル分け5は、学習が終わったら必要なら行ってください。\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "L-gAIubBAWvQ"
   },
   "source": [
    "## 0. 環境構築\n",
    "\n",
    "Style-Bert-VITS2の環境をcolab上に構築します。ランタイムがT4等のGPUバックエンドになっていることを確認し、実行してください。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "0GNj8JyDAlm2",
    "outputId": "d8be4a1a-e52d-46f8-8675-3f1a24bc9a51"
   },
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "os.environ[\"PATH\"] += \":/root/.cargo/bin\"\n",
    "\n",
    "!curl -LsSf https://astral.sh/uv/install.sh | sh\n",
    "!git clone https://github.com/litagin02/Style-Bert-VITS2.git\n",
    "%cd Style-Bert-VITS2/\n",
    "!uv pip install --system -q -r requirements-colab.txt\n",
    "!python initialize.py --skip_default_models"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "o5z1nzkvAWvR",
    "outputId": "cd87f053-18e0-4dbb-f904-d5230d1fa7ef"
   },
   "outputs": [],
   "source": [
    "# Google driveを使う方はこちらを実行してください。\n",
    "\n",
    "from google.colab import drive\n",
    "\n",
    "drive.mount(\"/content/drive\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "WU9apXzcAWvR"
   },
   "source": [
    "## 1. 初期設定\n",
    "\n",
    "学習とその結果を保存するディレクトリ名を指定します。\n",
    "Google driveの場合はそのまま実行、カスタマイズしたい方は変更して実行してください。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "gO3OwZV1AWvR"
   },
   "outputs": [],
   "source": [
    "# 学習に必要なファイルや途中経過が保存されるディレクトリ\n",
    "dataset_root = \"/content/drive/MyDrive/Style-Bert-VITS2/Data\"\n",
    "\n",
    "# 学習結果（音声合成に必要なファイルたち）が保存されるディレクトリ\n",
    "assets_root = \"/content/drive/MyDrive/Style-Bert-VITS2/model_assets\"\n",
    "\n",
    "import yaml\n",
    "\n",
    "\n",
    "with open(\"configs/paths.yml\", \"w\", encoding=\"utf-8\") as f:\n",
    "    yaml.dump({\"dataset_root\": dataset_root, \"assets_root\": assets_root}, f)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "dA_yLeezAWvS"
   },
   "source": [
    "## 2. 学習に使うデータ準備\n",
    "\n",
    "すでに音声ファイル（1ファイル2-12秒程度）とその書き起こしデータがある場合は2.2を、ない場合は2.1を実行してください。"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "8s9gOnTCAWvS"
   },
   "source": [
    "### 2.1 音声ファイルからのデータセットの作成（ある人はスキップ可）\n",
    "\n",
    "音声ファイル（1ファイル2-12秒程度）とその書き起こしのデータセットを持っていない方は、（日本語の）音声ファイルのみから以下の手順でデータセットを作成することができます。Google drive上の`Style-Bert-VITS2/inputs/`フォルダに音声ファイル（wavやmp3等の通常の音声ファイル形式、1ファイルでも複数ファイルでも可）を置いて、下を実行すると、データセットが作られ、自動的に正しい場所へ配置されます。\n",
    "\n",
    "**2024-06-02のVer 2.5以降**、`inputs/`フォルダにサブフォルダを2個以上作ってそこへ音声ファイルをスタイルに応じて振り分けて置くと、学習の際にサブディレクトリに応じたスタイルが自動的に作成されます。デフォルトスタイルのみでよい場合や手動でスタイルを後で作成する場合は`inputs/`直下へ入れれば大丈夫です。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "_fXCTPuiAWvS",
    "outputId": "47abd55b-efe5-48e2-f6fa-8e2016efe0ec"
   },
   "outputs": [],
   "source": [
    "# 元となる音声ファイル（wav形式）を入れるディレクトリ\n",
    "input_dir = \"/content/drive/MyDrive/Style-Bert-VITS2/inputs\"\n",
    "# モデル名（話者名）を入力\n",
    "model_name = \"your_model_name\"\n",
    "\n",
    "# こういうふうに書き起こして欲しいという例文（句読点の入れ方・笑い方や固有名詞等）\n",
    "initial_prompt = \"こんにちは。元気、ですかー？ふふっ、私は……ちゃんと元気だよ！\"\n",
    "\n",
    "!python slice.py -i {input_dir} --model_name {model_name}\n",
    "!python transcribe.py --model_name {model_name} --initial_prompt {initial_prompt} --use_hf_whisper"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "j7vEWewoAWvS"
   },
   "source": [
    "成功したらそのまま3へ進んでください"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "Z3AC-3zpAWvS"
   },
   "source": [
    "### 2.2 音声ファイルと書き起こしデータがすでにある場合\n",
    "\n",
    "指示に従って適切にデータセットを配置してください。\n",
    "\n",
    "次のセルを実行して、学習データをいれるフォルダ（1で設定した`dataset_root`）を作成します。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "esCNJl704h52"
   },
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "os.makedirs(dataset_root, exist_ok=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "aaDgJCjCAWvT"
   },
   "source": [
    "まず音声データと、書き起こしテキストを用意してください。\n",
    "\n",
    "それを次のように配置します。\n",
    "```\n",
    "├── Data/\n",
    "│   ├── {モデルの名前}\n",
    "│   │   ├── esd.list\n",
    "│   │   ├── raw/\n",
    "│   │   │   ├── foo.wav\n",
    "│   │   │   ├── bar.mp3\n",
    "│   │   │   ├── style1/\n",
    "│   │   │   │   ├── baz.wav\n",
    "│   │   │   │   ├── qux.wav\n",
    "│   │   │   ├── style2/\n",
    "│   │   │   │   ├── corge.wav\n",
    "│   │   │   │   ├── grault.wav\n",
    "...\n",
    "```\n",
    "\n",
    "### 配置の仕方\n",
    "- 上のように配置すると、`style1/`と`style2/`フォルダの内部（直下以外も含む）に入っている音声ファイルたちから、自動的にデフォルトスタイルに加えて`style1`と`style2`というスタイルが作成されます\n",
    "- 特にスタイルを作る必要がない場合や、スタイル分類機能等でスタイルを作る場合は、`raw/`フォルダ直下に全てを配置してください。このように`raw/`のサブディレクトリの個数が0または1の場合は、スタイルはデフォルトスタイルのみが作成されます。\n",
    "- 音声ファイルのフォーマットはwav形式以外にもmp3等の多くの音声ファイルに対応しています\n",
    "\n",
    "### 書き起こしファイル`esd.list`\n",
    "\n",
    "`Data/{モデルの名前}/esd.list` ファイルには、以下のフォーマットで各音声ファイルの情報を記述してください。\n",
    "\n",
    "\n",
    "```\n",
    "path/to/audio.wav(wavファイル以外でもこう書く)|{話者名}|{言語ID、ZHかJPかEN}|{書き起こしテキスト}\n",
    "```\n",
    "\n",
    "- ここで、最初の`path/to/audio.wav`は、`raw/`からの相対パスです。つまり、`raw/foo.wav`の場合は`foo.wav`、`raw/style1/bar.wav`の場合は`style1/bar.wav`となります。\n",
    "- 拡張子がwavでない場合でも、`esd.list`には`wav`と書いてください、つまり、`raw/bar.mp3`の場合でも`bar.wav`と書いてください。\n",
    "\n",
    "\n",
    "例：\n",
    "```\n",
    "foo.wav|hanako|JP|こんにちは、元気ですか？\n",
    "bar.wav|taro|JP|はい、聞こえています……。何か用ですか？\n",
    "style1/baz.wav|hanako|JP|今日はいい天気ですね。\n",
    "style1/qux.wav|taro|JP|はい、そうですね。\n",
    "...\n",
    "english_teacher.wav|Mary|EN|How are you? I'm fine, thank you, and you?\n",
    "...\n",
    "```\n",
    "もちろん日本語話者の単一話者データセットでも構いません。"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "5r85-W20ECcr"
   },
   "source": [
    "## 3. 学習の前処理\n",
    "\n",
    "次に学習の前処理を行います。必要なパラメータをここで指定します。次のセルに設定等を入力して実行してください。「～～かどうか」は`True`もしくは`False`を指定してください。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "CXR7kjuF5GlE"
   },
   "outputs": [],
   "source": [
    "# 上でつけたフォルダの名前`Data/{model_name}/`\n",
    "model_name = \"your_model_name\"\n",
    "\n",
    "# JP-Extra （日本語特化版）を使うかどうか。日本語の能力が向上する代わりに英語と中国語は使えなくなります。\n",
    "use_jp_extra = True\n",
    "\n",
    "# 学習のバッチサイズ。VRAMのはみ出具合に応じて調整してください。\n",
    "batch_size = 4\n",
    "\n",
    "# 学習のエポック数（データセットを合計何周するか）。\n",
    "# 100で多すぎるほどかもしれませんが、もっと多くやると質が上がるのかもしれません。\n",
    "epochs = 100\n",
    "\n",
    "# 保存頻度。何ステップごとにモデルを保存するか。分からなければデフォルトのままで。\n",
    "save_every_steps = 1000\n",
    "\n",
    "# 音声ファイルの音量を正規化するかどうか\n",
    "normalize = False\n",
    "\n",
    "# 音声ファイルの開始・終了にある無音区間を削除するかどうか\n",
    "trim = False\n",
    "\n",
    "# 読みのエラーが出た場合にどうするか。\n",
    "# \"raise\"ならテキスト前処理が終わったら中断、\"skip\"なら読めない行は学習に使わない、\"use\"なら無理やり使う\n",
    "yomi_error = \"skip\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "BFZdLTtpAWvT"
   },
   "source": [
    "上のセルが実行されたら、次のセルを実行して学習の前処理を行います。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "xMVaOIPLabV5",
    "outputId": "36b1c2b2-6df0-4d00-d86a-519a0fc0af63"
   },
   "outputs": [],
   "source": [
    "from gradio_tabs.train import preprocess_all\n",
    "from style_bert_vits2.nlp.japanese import pyopenjtalk_worker\n",
    "\n",
    "\n",
    "pyopenjtalk_worker.initialize_worker()\n",
    "\n",
    "preprocess_all(\n",
    "    model_name=model_name,\n",
    "    batch_size=batch_size,\n",
    "    epochs=epochs,\n",
    "    save_every_steps=save_every_steps,\n",
    "    num_processes=2,\n",
    "    normalize=normalize,\n",
    "    trim=trim,\n",
    "    freeze_EN_bert=False,\n",
    "    freeze_JP_bert=False,\n",
    "    freeze_ZH_bert=False,\n",
    "    freeze_style=False,\n",
    "    freeze_decoder=False,\n",
    "    use_jp_extra=use_jp_extra,\n",
    "    val_per_lang=0,\n",
    "    log_interval=200,\n",
    "    yomi_error=yomi_error,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "sVhwI5C-AWvT"
   },
   "source": [
    "## 4. 学習\n",
    "\n",
    "前処理が正常に終わったら、学習を行います。次のセルを実行すると学習が始まります。\n",
    "\n",
    "学習の結果は、上で指定した`save_every_steps`の間隔で、Google Driveの中の`Style-Bert-VITS2/Data/{モデルの名前}/model_assets/`フォルダに保存されます。\n",
    "\n",
    "このフォルダをダウンロードし、ローカルのStyle-Bert-VITS2の`model_assets`フォルダに上書きすれば、学習結果を使うことができます。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "laieKrbEb6Ij"
   },
   "outputs": [],
   "source": [
    "# 上でつけたモデル名を入力。学習を途中からする場合はきちんとモデルが保存されているフォルダ名を入力。\n",
    "model_name = \"your_model_name\"\n",
    "\n",
    "\n",
    "import yaml\n",
    "from gradio_tabs.train import get_path\n",
    "\n",
    "paths = get_path(model_name)\n",
    "dataset_path = str(paths.dataset_path)\n",
    "config_path = str(paths.config_path)\n",
    "\n",
    "with open(\"default_config.yml\", \"r\", encoding=\"utf-8\") as f:\n",
    "    yml_data = yaml.safe_load(f)\n",
    "yml_data[\"model_name\"] = model_name\n",
    "with open(\"config.yml\", \"w\", encoding=\"utf-8\") as f:\n",
    "    yaml.dump(yml_data, f, allow_unicode=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "background_save": true,
     "base_uri": "https://localhost:8080/"
    },
    "id": "JqGeHNabAWvT",
    "outputId": "c51b422c-728b-420b-fa92-b787fa058adf"
   },
   "outputs": [],
   "source": [
    "# 日本語特化版を「使う」場合\n",
    "!python train_ms_jp_extra.py --config {config_path} --model {dataset_path} --assets_root {assets_root}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "rVbjh-WPAWvU"
   },
   "outputs": [],
   "source": [
    "# 日本語特化版を「使わない」場合\n",
    "!python train_ms.py --config {config_path} --model {dataset_path} --assets_root {assets_root}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "c7g0hrdeP1Tl",
    "outputId": "4bb9d21e-50df-4ba5-a547-daa78a4b63dc"
   },
   "outputs": [],
   "source": [
    "# 学習結果を試す・マージ・スタイル分けはこちらから\n",
    "!python app.py --share"
   ]
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "gpuType": "T4",
   "provenance": []
  },
  "kernelspec": {
   "display_name": "Python 3",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}
